#include <uv.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <cstring>
#include <string>
#include <unordered_map>

// ----------------------------
//          TLS 封装
// ----------------------------
class TLSWrapper {
public:
    SSL_CTX* ssl_ctx;
    SSL* ssl;
    BIO* bio_read;
    BIO* bio_write;

    TLSWrapper(SSL_CTX* ctx) : ssl_ctx(ctx), ssl(nullptr) {
        bio_read = BIO_new(BIO_s_mem());
        bio_write = BIO_new(BIO_s_mem());
    }

    ~TLSWrapper() {
        if (ssl) SSL_free(ssl);
        BIO_free(bio_read);
        BIO_free(bio_write);
    }

    bool handshake() {
        if (!ssl) {
            ssl = SSL_new(ssl_ctx);
            SSL_set_bio(ssl, bio_read, bio_write);
            SSL_set_accept_state(ssl);
        }
        return SSL_do_handshake(ssl) == 1;
    }
};

// ----------------------------
//      WebSocket 连接管理
// ----------------------------
class WebSocketConnection {
public:
    uv_tcp_t* handle;
    TLSWrapper* tls;
    std::string buffer;

    WebSocketConnection(uv_tcp_t* h) : handle(h), tls(nullptr) {}
};

// ----------------------------
//        全局变量
// ----------------------------
std::unordered_map<uv_tcp_t*, WebSocketConnection*> connections;
SSL_CTX* ssl_ctx;

// ----------------------------
//        TLS 初始化
// ----------------------------
void init_openssl() {
    SSL_load_error_strings();
    OpenSSL_add_ssl_algorithms();
    ssl_ctx = SSL_CTX_new(TLS_server_method());
    SSL_CTX_use_certificate_file(ssl_ctx, "server.crt", SSL_FILETYPE_PEM);
    SSL_CTX_use_PrivateKey_file(ssl_ctx, "server.key", SSL_FILETYPE_PEM);
}

// ----------------------------
//      WebSocket 握手响应
// ----------------------------
std::string get_ws_accept(const std::string& key) {
    // 实现 WebSocket Accept 密钥生成逻辑
    return "HTTP/1.1 101 Switching Protocols\r\n"
           "Upgrade: websocket\r\n"
           "Connection: Upgrade\r\n"
           "Sec-WebSocket-Accept: ...\r\n\r\n";
}

// ----------------------------
//       libuv 回调函数
// ----------------------------
void alloc_buffer(uv_handle_t* handle, size_t size, uv_buf_t* buf) {
    buf->base = new char[size];
    buf->len = size;
}

void on_close(uv_handle_t* handle) {
    auto conn = connections[(uv_tcp_t*)handle];
    delete conn;
    connections.erase((uv_tcp_t*)handle);
}

void on_write(uv_write_t* req, int status) {
    delete req;
}

void on_read(uv_stream_t* stream, ssize_t nread, const uv_buf_t* buf) {
    auto conn = connections[(uv_tcp_t*)stream];
    if (nread > 0) {
        // 处理 TLS 解密
        BIO_write(conn->tls->bio_read, buf->base, nread);
        if (!conn->tls->handshake()) {
            // TLS 握手未完成
            return;
        }

        // 解密后的数据
        char decrypted[4096];
        int decrypted_len = SSL_read(conn->tls->ssl, decrypted, sizeof(decrypted));

        // WebSocket 握手
        if (conn->buffer.find("Upgrade: websocket") != std::string::npos) {
            std::string response = get_ws_accept("...");
            uv_write_t* req = new uv_write_t;
            uv_buf_t res_buf = uv_buf_init((char*)response.c_str(), response.size());
            uv_write(req, stream, &res_buf, 1, on_write);
        }

        // 处理 WebSocket 数据帧
        // ... (需实现 WebSocket 协议解析)
    }
    delete[] buf->base;
}

void on_connection(uv_stream_t* server, int status) {
    uv_tcp_t* client = new uv_tcp_t;
    uv_tcp_init(uv_default_loop(), client);
    uv_accept(server, (uv_stream_t*)client);

    auto conn = new WebSocketConnection(client);
    conn->tls = new TLSWrapper(ssl_ctx);
    connections[client] = conn;

    uv_read_start((uv_stream_t*)client, alloc_buffer, on_read);
}

// ----------------------------
//          主函数
// ----------------------------
int main() {
    // 初始化 OpenSSL
    init_openssl();

    // 初始化 libuv
    uv_loop_t* loop = uv_default_loop();
    uv_tcp_t server;
    uv_tcp_init(loop, &server);

    // 绑定端口
    struct sockaddr_in addr;
    uv_ip4_addr("0.0.0.0", 8443, &addr);
    uv_tcp_bind(&server, (const struct sockaddr*)&addr, 0);

    // 开始监听
    uv_listen((uv_stream_t*)&server, 128, on_connection);

    // 启动事件循环
    uv_run(loop, UV_RUN_DEFAULT);

    // 清理
    SSL_CTX_free(ssl_ctx);
    EVP_cleanup();
    return 0;
}
// # 安装依赖
// sudo apt-get install libuv1-dev libssl-dev

// # 编译命令
// g++ -std=c++11 ws_server.cpp -o ws_server -luv -lssl -lcrypto
// openssl req -x509 -newkey rsa:4096 -keyout server.key -out server.crt -days 365 -nodes
